#!/usr/bin/env python3
# encoding: utf-8
# @data:2023/07/03
# @author:aiden
# 视觉巡线节点
import cv2
import math
import queue
import rospy
import signal
import numpy as np
from ainex_sdk import common
from sensor_msgs.msg import Image
from ainex_kinematics.gait_manager import GaitManager
from ainex_kinematics.motion_manager import MotionManager

LAB_CONFIG_PATH = '/home/ubuntu/software/lab_tool/lab_config.yaml'
class VisualPatrolNode():
    def __init__(self, name):
        rospy.init_node(name)
        self.name = name
        self.running = True
        self.count_finish = 0
        self.image_queue = queue.Queue(maxsize=2)
        self.motion_manager = MotionManager()
        lab_config = common.get_yaml_data(LAB_CONFIG_PATH)
        self.lower = tuple(lab_config['lab']['Mono']['white']['min'])
        self.upper = tuple(lab_config['lab']['Mono']['white']['max'])
        self.gait_manager = GaitManager()
        self.go_gait_param = self.gait_manager.get_gait_param()

        self.go_gait_param['init_x_offset'] = 0
        self.go_gait_param['init_y_offset'] = 0
        self.go_gait_param['init_roll_offset'] = 0
        self.go_gait_param['init_pitch_offset'] = 0
        self.go_gait_param['init_yaw_offset'] = 0

        self.go_gait_param['body_height'] = 0.035
        self.go_gait_param['step_height'] = 0.025
        self.go_gait_param['hip_pitch_offset'] = 15
        self.go_gait_param['z_swap_amplitude'] = 0
        self.go_dsp = [100, 0.1, 0.01]
        
        self.go_arm_swap = 30
        self.camera = rospy.get_param('/camera')
        rospy.Subscriber('/{}/{}'.format(self.camera['camera_name'], self.camera['image_topic']), Image, self.image_callback)  # 订阅图像
        # 设置退出处理函数
        signal.signal(signal.SIGINT, self.shutdown)

        self.motion_manager.run_action('walk_ready')     # 播放准备姿势动作
  
    def get_area_max_contour(self, contours, min_area=50, max_area=640*480):
        """
        获取轮廓中面积最重大的一个, 过滤掉面积过小的情况(get the contour whose area is the largest. Filter out those whose area is too small)
        :param contours: 轮廓列表(contour list)
        :param threshold: 面积阈值, 小于这个面积的轮廓会被过滤(area threshold. Contour whose area is less than this value will be filtered out)
        :return: 如果最大的轮廓面积大于阈值则返回最大的轮廓, 否则返回None(if the maximum contour area is greater than this threshold, return the
        largest contour, otherwise return None)
        """
        contour_area_max = 0
        area_max_contour = None

        for c in contours:  # 历遍所有轮廓
            contour_area_temp = math.fabs(cv2.contourArea(c))  # 计算轮廓面积
            if contour_area_temp > contour_area_max:
                contour_area_max = contour_area_temp
                if max_area >= contour_area_temp >= min_area:  # 过滤干扰
                    area_max_contour = c

        return area_max_contour,  contour_area_max  # 返回最大的

    def image_callback(self, ros_image):
        rgb_image = np.ndarray(shape=(ros_image.height, ros_image.width, 3), dtype=np.uint8,
                           buffer=ros_image.data)  # 将ros格式图像消息转化为opencv格式
        bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
        image_resize = cv2.resize(bgr_image, (320, 240), interpolation=cv2.INTER_NEAREST)  # 图像缩放, 加快图像处理速度, 不能太小，否则会丢失大量细节
        image_gb = cv2.GaussianBlur(image_resize, (3, 3), 3)
        image_lab = cv2.cvtColor(image_gb, cv2.COLOR_BGR2LAB)  # bgr空间转lab空间，方便提取颜色

        binary = cv2.inRange(image_lab, self.lower, self.upper)  # 二值化
        element = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        eroded = cv2.erode(binary, element)  # 腐蚀
        dilated = cv2.dilate(eroded, cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)))  # 膨胀
        roi_left = dilated[:, 0:160]
        roi_right = dilated[:, 160:320]

        contours_left = cv2.findContours(roi_left, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[-2]  # 找出所有轮廓
        contours_right = cv2.findContours(roi_right, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[-2]  # 找出所有轮廓
        left_point = -1
        right_point = -1 
        
        max_contour_left, contour_area = self.get_area_max_contour(contours_left, 10, 160*120)
        if max_contour_left is not None:
            rect = cv2.minAreaRect(max_contour_left)  # 获取最小外接矩形
            corners = np.int64(cv2.boxPoints(rect))  # 获取最小外接矩形的四个角点
            for j in range(4):
                corners[j, 0] = 2*corners[j, 0]
                corners[j, 1] = 2*corners[j, 1]
            cv2.drawContours(bgr_image, [corners], -1, (0, 255, 255), 2)  # 画出四个点组成的矩形
            bottom_point = corners[corners[:, 1].argmax()]
            cv2.circle(bgr_image, tuple(bottom_point), 5, (0, 0, 255), -1)
            left_point = bottom_point[1]
        
        max_contour_right, contour_area = self.get_area_max_contour(contours_right, 10, 160*120)
        if max_contour_right is not None:
            rect = cv2.minAreaRect(max_contour_right)
            corners = np.int64(cv2.boxPoints(rect))  # 获取最小外接矩形的四个角点
            for j in range(4):
                corners[j, 0] = 2*corners[j, 0] + 320
                corners[j, 1] = 2*corners[j, 1]
            cv2.drawContours(bgr_image, [corners], -1, (0, 255, 255), 2)  # 画出四个点组成的矩形
            bottom_point = corners[corners[:, 1].argmax()]
            cv2.circle(bgr_image, tuple(bottom_point), 5, (0, 0, 255), -1)
            right_point = bottom_point[1]
        
        if not self.image_queue.empty():
            try:
                self.image_queue.get_nowait()
            except queue.Empty:
                pass
        try:
            self.image_queue.put_nowait([bgr_image, left_point, right_point])
        except queue.Full:
            pass

    # 节点关闭回调函数
    def shutdown(self, signum, frame):
        self.motion_manager.run_action('stand')
        self.running = False 
        common.loginfo('%s shutdown' % self.name)

    # 主循环函数
    def run(self):
        while self.running:
            image, left_point, right_point = self.image_queue.get(block=True)
            # print(left_point, right_point)
            if 1:
                if left_point != -1:
                    self.count_finish = 0
                    if right_point != -1:
                        if left_point > right_point:
                            if left_point - right_point > 50:
                                self.gait_manager.set_step(self.go_dsp, 0.02, 0, -5, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                            else:
                                self.gait_manager.set_step(self.go_dsp, 0.03, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                        else:
                            if left_point - right_point < -50:
                                self.gait_manager.set_step(self.go_dsp, 0.02, 0, 5, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                            else:
                                self.gait_manager.set_step(self.go_dsp, 0.03, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                    else:
                        self.gait_manager.set_step(self.go_dsp, 0.02, 0, -5, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                elif right_point != -1:
                    self.count_finish = 0
                    self.gait_manager.set_step(self.go_dsp, 0.02, 0, 5, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                else:
                    self.count_finish += 1
                    self.gait_manager.set_step(self.go_dsp, 0.03, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                if self.count_finish > 20:
                    print('finish')
                    break
            cv2.imshow('image', image)
            key = cv2.waitKey(1)
            if key != -1:
                break
        self.gait_manager.stop()
        rospy.signal_shutdown('shutdown')

if __name__ == "__main__":
    VisualPatrolNode('sprint_1').run()
